/**
 * Copyright 2011 Brigham Young University
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package edu.byu.nlp.data;

import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;

import org.apache.commons.math3.linear.AbstractRealMatrix;
import org.apache.commons.math3.linear.MatrixUtils;
import org.apache.commons.math3.linear.RealMatrix;
import org.apache.commons.math3.linear.RealMatrixPreservingVisitor;

import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;

/**
 * A utility class for building a feature matrix that consists solely of block features. In this case, the matrix can
 * actually be efficiently represented by a feature vector. Each entry in the underlying feature vector gets expanded to
 * a square matrix whose value is the value in the feature vector times the identity matrix. If {@code getEntry()} is to
 * be called on the returned {@code FeatureMatrix}, then the {@code sortOnBuild} flag must be set upon construction,
 * otherwise {@code getEntry()} will not operate correctly. Also note the {@code sortOnBuild} take time logarithmic in
 * the number of active (non-zero) features.
 * 
 * @author rah67
 *
 */
public class BlockFeatureMatrixBuilder {
	
	private static class FeatureValuePair {
		private final int featureIndex;
		private final double value;
		
		public FeatureValuePair(int featureIndex, double value) {
			this.featureIndex = featureIndex;
			this.value = value;
		}

		public int getFeatureIndex() {
			return featureIndex;
		}

		public double getValue() {
			return value;
		}
	}
	
	private static class FeatureIndexComparator implements Comparator<FeatureValuePair> {
		@Override
		public int compare(FeatureValuePair pair1, FeatureValuePair pair2) {
			return pair1.featureIndex - pair2.featureIndex;
		}
	}
	
	private final boolean sortOnBuild;
	private final List<FeatureValuePair> pairs;
	private final Dimension dimension;
	
	/**
	 * @param dimension the number of labels by the number of features (i.e. columns; non-block features * labels)
	 * @param sortOnBuild indicates whether feature vectors are sorted when {@code #buildFeatureMatrix()} is called
	 */
	public BlockFeatureMatrixBuilder(Dimension dimension, boolean sortOnBuild) {
		this.dimension = dimension;
		this.sortOnBuild = sortOnBuild;
		this.pairs = Lists.newArrayList();
	}
	
	public boolean addFeatureValuePair(int featureIndex, double value) {
		return pairs.add(new FeatureValuePair(featureIndex, value));
	}
	
	public void clear() {
		pairs.clear();
	}
	
	private static class BlockFeatureMatrix extends AbstractRealMatrix implements FeatureMatrix {
		
		private final Dimension indexInfo;
		private final int[] featureIndices;
		private final double[] values;

		/**
		 * featureIndices must be sorted for the getEntry() to operate correctly (unchecked).
		 * 
		 * @throws NullPointerException if any of the objects are null.
		 * @throws IllegalArgumentException if featureIndices.length != values.length
		 */
		public BlockFeatureMatrix(Dimension indexInfo, int[] featureIndices, double[] values) {
			Preconditions.checkNotNull(indexInfo);
			Preconditions.checkNotNull(featureIndices);
			Preconditions.checkNotNull(values);
			Preconditions.checkArgument(featureIndices.length == values.length);
			this.indexInfo = indexInfo;
			this.featureIndices = featureIndices;
			this.values = values;
		}

		/** {@inheritDoc} */
		@Override
		public void walkSparselyInOptimizedOrder(RealMatrixPreservingVisitor visitor) {
			// If we walk by columns, we will hit a single block feature repeatedly for each row which will give better
			// cache performance.
			walkSparselyInColumnOrder(visitor);
		}

		/** {@inheritDoc} */
		@Override
		public void walkSparselyInRowOrder(RealMatrixPreservingVisitor visitor) {
			// TODO Auto-generated method stub
		}

		/** {@inheritDoc} */
		@Override
		public void walkSparselyInColumnOrder(RealMatrixPreservingVisitor visitor) {
			for (int i = 0; i < featureIndices.length; i++) {
				// A block feature is equivalent to a square submatrix whose value is the value of the block feature
				// times the identity matrix. This inner loop traverses the non-zero diagonal.

				// The start of the square sub-matrix
				int column = featureIndices[i] * indexInfo.getNumLabels();
				for (int label = 0; label < indexInfo.getNumLabels(); label++, column++) {
					visitor.visit(column, label, values[i]);
				}
			}
		}

		/** {@inheritDoc} */
		@Override
		public void addToEntry(int label, int columnIndex, double value) {
			throw new UnsupportedOperationException("Cannot add to indiviual entries in a block feature vector are" +
					"immutable");
		}

		/** {@inheritDoc} */
		@Override
		public RealMatrix copy() {
			int[] featureIndices = this.featureIndices.clone();
			double[] values = this.values.clone();
			return new BlockFeatureMatrix(this.indexInfo, featureIndices, values);
		}

		/**
		 * {@inheritDoc}
		 */
		@Override
		public RealMatrix createMatrix(int rowDimension, int colDimension) throws IllegalArgumentException {
			throw new UnsupportedOperationException("Creating a new matrix is pointless since it is immutable");
		}

		/**
		 * {@inheritDoc}
		 */
		@Override
		public int getColumnDimension() {
			return indexInfo.getNumFeatures() * indexInfo.getNumLabels();
		}

		/**
		 * {@inheritDoc}
		 * 
		 * This is not a constant time operation. Requires time logarithmic in the number of non-zero entries in the
		 * block feature vector. 
		 */
		@Override
		public double getEntry(int label, int nonBlockFeatureIndex) {
			MatrixUtils.checkRowIndex(this, label);
			MatrixUtils.checkColumnIndex(this, nonBlockFeatureIndex);

			// We divide by the number of classes to find the feature index in the block feature vector
			int index = Arrays.binarySearch(featureIndices, nonBlockFeatureIndex / indexInfo.getNumLabels());
			if (index < 0) {
				return 0.0;
			}
			return values[index];
		}

		/** {@inheritDoc} */
		@Override
		public int getRowDimension() {
			return indexInfo.getNumLabels();
		}

		/** {@inheritDoc} */
		@Override
		public void multiplyEntry(int label, int nonBlockFeatureIndex, double value) {
			throw new UnsupportedOperationException("Cannot multiply to indiviual entries in a block feature vector " +
					"are immutable");
		}

		/** {@inheritDoc} */
		@Override
		public void setEntry(int arg0, int arg1, double arg2) {
			throw new UnsupportedOperationException("Cannot set indiviual entries in a block feature vector are " +
					"immutable");
		}
		
	}
	
	public FeatureMatrix buildFeatureMatrix() {
		if (sortOnBuild) {
			Collections.sort(pairs, new FeatureIndexComparator());
		}
		int[] featureIndices = extractIndices();
		double[] values = extractValues();
		return new BlockFeatureMatrix(dimension, featureIndices, values);
	}
	
	private int[] extractIndices() {
		int[] featureIndices = new int[pairs.size()];
		for (int i = 0; i < featureIndices.length; i++) {
			featureIndices[i] = pairs.get(i).getFeatureIndex();
		}
		return featureIndices;
	}

	private double[] extractValues() {
		double[] values = new double[pairs.size()];
		for (int i = 0; i < values.length; i++) {
			values[i] = pairs.get(i).getValue();
		}
		return values;
	}
}
